from typing import Dict, Any, List, Union
import json
from pathlib import Path
import numpy as np
import cv2


class PDFExtractor:
    """
    PDF提取器，使用pypdfium2处理PDF文件，支持按页面提取内容并输出多种格式
    """

    def __init__(self):
        """
        初始化PDF提取器
        """
        # 检查pypdfium2是否可用
        try:
            import pypdfium2 as pdfium
            self.pdfium = pdfium
            self.available = True
        except ImportError:
            self.pdfium = None
            self.available = False

    def is_available(self) -> bool:
        """
        检查PDF提取器是否可用

        Returns:
            bool: 如果pypdfium2可用返回True，否则返回False
        """
        return self.available

    def extract_pages_as_images(self, pdf_path: str, rotation: int = 0, zoom: float = 2.0) -> List[np.ndarray]:
        """
        从PDF文件中提取每页作为图像

        Args:
            pdf_path (str): PDF文件路径
            rotation (int): 旋转角度，默认为0
            zoom (float): 缩放比例，默认为2.0
                       基于96 DPI基准，换算关系为：DPI = 96 * zoom
                       常用设置：
                       - 屏幕显示：zoom=1.0-2.0（96-192 DPI）
                       - OCR识别：zoom=2.0-3.0（192-288 DPI）
                       - 打印/出版：zoom=3.125（300 DPI）

        Returns:
            List[np.ndarray]: 每页的图像数据列表
        """
        if not self.is_available():
            raise ImportError("pypdfium2 is not available. Please install it to use this feature.")

        images = []
        # 打开PDF文档
        doc = self.pdfium.PdfDocument(pdf_path)
        try:
            # 遍历每一页
            for page in doc:
                # 将页面渲染为图像
                image = page.render(rotation=rotation, scale=zoom).to_pil()
                # 转换为RGB模式确保颜色一致性
                image = image.convert("RGB")
                # 转换为OpenCV格式（BGR）
                img_cv = np.array(image)
                img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
                images.append(img_cv)
        finally:
            # 确保在使用完毕后关闭PDF文档
            doc.close()

        return images

    def extract_text_from_page(self, pdf_path: str, page_number: int) -> str:
        """
        从PDF的指定页面提取文本

        Args:
            pdf_path (str): PDF文件路径
            page_number (int): 页码（从0开始）

        Returns:
            str: 提取的文本内容
        """
        if not self.is_available():
            raise ImportError("pypdfium2 is not available. Please install it to use this feature.")

        doc = self.pdfium.PdfDocument(pdf_path)
        try:
            page = doc[page_number]
            text_page = page.get_textpage()
            return text_page.get_text_bounded()
        finally:
            doc.close()

    def extract_text_from_all_pages(self, pdf_path: str) -> List[str]:
        """
        从PDF的所有页面提取文本

        Args:
            pdf_path (str): PDF文件路径

        Returns:
            List[str]: 每页的文本内容列表
        """
        if not self.is_available():
            raise ImportError("pypdfium2 is not available. Please install it to use this feature.")

        texts = []
        doc = self.pdfium.PdfDocument(pdf_path)
        try:
            for page in doc:
                text_page = page.get_textpage()
                texts.append(text_page.get_text_bounded())
        finally:
            doc.close()

        return texts

    def get_page_count(self, pdf_path: str) -> int:
        """
        获取PDF文件的总页数

        Args:
            pdf_path (str): PDF文件路径

        Returns:
            int: PDF文件的总页数
        """
        if not self.is_available():
            raise ImportError("pypdfium2 is not available. Please install it to use this feature.")

        doc = self.pdfium.PdfDocument(pdf_path)
        try:
            return len(doc)
        finally:
            doc.close()

    def extract_to_images(self, pdf_path: str, output_dir: str,
                          image_format: str = "jpg",
                          rotation: int = 0,
                          zoom: float = 2.0) -> List[str]:
        """
        将PDF每页提取为图像文件并保存

        Args:
            pdf_path (str): PDF文件路径
            output_dir (str): 输出目录路径
            image_format (str): 图像格式（jpg, png），默认为jpg
            rotation (int): 旋转角度，默认为0
            zoom (float): 缩放比例，默认为2.0

        Returns:
            List[str]: 保存的图像文件路径列表
        """
        if not self.is_available():
            raise ImportError("pypdfium2 is not available. Please install it to use this feature.")

        # 创建输出目录
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)

        # 提取图像
        images = self.extract_pages_as_images(pdf_path, rotation, zoom)

        # 保存图像文件
        saved_files = []
        for i, img in enumerate(images):
            filename = f"page_{i + 1:03d}.{image_format}"
            filepath = str(output_path / filename)

            if image_format.lower() == "png":
                cv2.imwrite(filepath, img, [cv2.IMWRITE_PNG_COMPRESSION, 9])
            else:  # 默认为jpg
                cv2.imwrite(filepath, img, [cv2.IMWRITE_JPEG_QUALITY, 95])

            saved_files.append(filepath)

        return saved_files

    def extract_to_text_files(self, pdf_path: str, output_dir: str) -> List[str]:
        """
        将PDF每页提取为文本文件并保存

        Args:
            pdf_path (str): PDF文件路径
            output_dir (str): 输出目录路径

        Returns:
            List[str]: 保存的文本文件路径列表
        """
        if not self.is_available():
            raise ImportError("pypdfium2 is not available. Please install it to use this feature.")

        # 创建输出目录
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)

        # 提取所有页面的文本
        texts = self.extract_text_from_all_pages(pdf_path)

        # 保存文本文件
        saved_files = []
        for i, text in enumerate(texts):
            filename = f"page_{i + 1:03d}.txt"
            filepath = str(output_path / filename)

            with open(filepath, "w", encoding="utf-8") as f:
                f.write(text)

            saved_files.append(filepath)

        return saved_files

    def extract_to_json(self, pdf_path: str, output_path: str) -> str:
        """
        将PDF内容提取为JSON格式并保存

        Args:
            pdf_path (str): PDF文件路径
            output_path (str): 输出JSON文件路径

        Returns:
            str: 保存的JSON文件路径
        """
        if not self.is_available():
            raise ImportError("pypdfium2 is not available. Please install it to use this feature.")

        # 获取PDF信息
        doc = self.pdfium.PdfDocument(pdf_path)
        try:
            version = doc.get_version()
            page_count = len(doc)
        finally:
            doc.close()

        # 提取所有页面的文本
        texts = self.extract_text_from_all_pages(pdf_path)

        # 构建JSON数据
        pdf_data = {
            "pdf_path": pdf_path,
            "pdf_version": version,
            "page_count": page_count,
            "pages": []
        }

        for i, text in enumerate(texts):
            page_data = {
                "page_number": i + 1,
                "text_content": text
            }
            pdf_data["pages"].append(page_data)

        # 保存JSON文件
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(pdf_data, f, ensure_ascii=False, indent=2)

        return output_path

    def extract_page_info(self, pdf_path: str, page_number: int) -> Dict[str, Any]:
        """
        获取PDF指定页面的详细信息

        Args:
            pdf_path (str): PDF文件路径
            page_number (int): 页码（从0开始）

        Returns:
            Dict[str, Any]: 页面信息字典
        """
        if not self.is_available():
            raise ImportError("pypdfium2 is not available. Please install it to use this feature.")

        doc = self.pdfium.PdfDocument(pdf_path)
        try:
            page = doc[page_number]
            width = page.get_width()
            height = page.get_height()

            # 获取文本内容
            text_page = page.get_textpage()
            text = text_page.get_text_bounded()

            return {
                "page_number": page_number + 1,
                "width": width,
                "height": height,
                "text_content": text,
                "word_count": len(text.split())
            }
        finally:
            doc.close()

    def extract_text_from_area(self, pdf_path: str, page_number: int, bbox: tuple) -> str:
        """
        从PDF指定页面的指定区域提取文本

        Args:
            pdf_path (str): PDF文件路径
            page_number (int): 页码（从0开始）
            bbox (tuple): 区域坐标 (x0, y0, x1, y1)

        Returns:
            str: 提取的文本内容
        """
        if not self.is_available():
            raise ImportError("pypdfium2 is not available. Please install it to use this feature.")

        doc = self.pdfium.PdfDocument(pdf_path)
        try:
            page = doc[page_number]
            text_page = page.get_textpage()
            # 根据bbox参数提取指定区域的文本
            # 注意：pypdfium2的坐标系统可能需要转换
            left, top, right, bottom = bbox
            return text_page.get_text_bounded(left=left, top=top, right=right, bottom=bottom)
        finally:
            doc.close()
